import asyncio
import concurrent.futures
import datetime
import functools
import json
import logging
import os
import signal
import subprocess
import sys
import tempfile
import time
from collections import defaultdict

from absl import app, flags
from tornado.ioloop import IOLoop, PeriodicCallback

from coin.base.logging import init_logging
from coin.base.timestamp import get_timestamp
from coin.tool.strat_monitor.util.info_util import _get_topics
from coin.util.queue.constants import TopicType
from coin.util.queue.tools.kafka_topic import parse_kafka_topic
from coin.tool.strat_monitor.util.monitor_util import separate_into_group


FLAGS = flags.FLAGS


def _write_strategy_file(strats, strat_file):
  json_str = json.dumps([{'strategy_name': strat} for strat in strats], indent=2)
  with open(strat_file, 'w') as fd:
    fd.write(json_str)


def _gen_flags(args):
  output = ' '
  for key, value in args.items():
    if isinstance(value, bool):
      if value is True:
        output += ' --%s' % key
    else:
      if value is not None:
        output += ' --%s %s' % (key, str(value))
  return output


def _get_strats_from_topic(topics):
  return list(set([parse_kafka_topic(topic)[1].strategy_name for topic in topics]))


def _to_td_str(ts):
  td = datetime.datetime.utcfromtimestamp(ts / 10**9)
  td_str = td.date().strftime('%Y%m%d')
  return td_str


class MonitorRunner(object):
  def __init__(self, *, ioloop, strategy_config_filename, max_workers, coin_env,
               business_unit=None, **args):
    self._start_ts = get_timestamp()
    self._td_str = _to_td_str(self._start_ts)
    self._ioloop = ioloop
    self._strategy_config_filename = strategy_config_filename
    kafka_config_filename = args.get('kafka_config_filename')
    self._topic_type = TopicType[args.get('kafka_topic_type')]
    self._exit_after_min = args.get('exit_after_min')
    assert kafka_config_filename is not None
    self._kafka_config_filename = kafka_config_filename
    self._monitor_args = args
    self._topics = _get_topics(
        self._strategy_config_filename, self._kafka_config_filename,
        topic_type=self._topic_type, business_unit=business_unit)
    self._strats = _get_strats_from_topic(self._topics)
    self._business_unit = business_unit
    self._max_workers = max_workers
    self._coin_env = coin_env
    self._monitor_info = {}
    self._strat_monitor_id_dict = {}
    self._new_monitor_id = None
    self._exit = False
    self._start()

  def _update_monitor_info(self, monitor_id, strats, strat_file, process):
    self._monitor_info[monitor_id] = {
        'strats': strats, 'strat_file': strat_file, 'process': process}

  def _update_strat_monitor_id_dict(self, strats, monitor_id):
    for strat in strats:
      self._strat_monitor_id_dict[strat] = monitor_id

  def _kill_monitor(self, monitor_id):
    self._monitor_info[monitor_id]['process'].terminate()

  def _start(self):
    PeriodicCallback(self._check_new_topics, 30*1000).start()
    PeriodicCallback(self._check_running_status, 60*1000).start()
    PeriodicCallback(self._check_exit, 10*1000).start()
    strat_groups = separate_into_group(self._strats, self._max_workers)
    for monitor_id, strats in enumerate(strat_groups):
      assert monitor_id not in self._monitor_info, monitor_id
      for strat in strats:
        assert strat not in self._strat_monitor_id_dict, strat
      self._launch_sub_monitor(monitor_id, strats)
      time.sleep(0.1)

  def _launch_sub_monitor(self, monitor_id, strats):
    assert len(strats) > 0
    if monitor_id in self._monitor_info:
      strat_file = self._monitor_info[monitor_id]['strat_file']
    else:
      with tempfile.NamedTemporaryFile(mode='w', dir='', delete=False) as fd:
        strat_file = fd.name
    if monitor_id in self._monitor_info:
      self._kill_monitor(monitor_id)
    _write_strategy_file(strats, strat_file)
    args = {'strategy_config_filename': strat_file}
    args.update(self._monitor_args)
    cmd = 'source ${HOME}/miniconda3/etc/profile.d/conda.sh' + ' && ' + \
          f'conda activate {self._coin_env}' + ' && ' + \
          './pyrunner coin/tool/strat_monitor/stream_processor_2.py' + \
          _gen_flags(args)
    monitor_process = subprocess.Popen(cmd, shell=True)
    self._update_monitor_info(monitor_id, strats, strat_file, monitor_process)
    self._update_strat_monitor_id_dict(strats, monitor_id)

  def _handle_new_strat_topic(self, topics):
    if len(topics) == 0:
      return
    new_strats = _get_strats_from_topic(topics)
    new_worker_needed = True
    if self._new_monitor_id is not None:
      strat_num = len(self._monitor_info[self._new_monitor_id]['strats'])
      avg_strats_per_worker = len(self._strats) / len(self._monitor_info)
      if strat_num < avg_strats_per_worker:
        new_worker_needed = False
    strats = new_strats
    if new_worker_needed:
      monitor_id = max(self._monitor_info.keys()) + 1
      logging.info('Launch new monitor process for new strat: %s' % new_strats)
    else:
      logging.info('Reuse existing monitor process for new strats: %s' % new_strats)
      assert self._new_monitor_id is not None
      monitor_id = self._new_monitor_id
      strats += self._monitor_info[monitor_id]['strats']
    self._new_monitor_id = monitor_id
    self._launch_sub_monitor(monitor_id, strats)

  def _handle_old_strat_topic(self, topics):
    if len(topics) == 0:
      return
    old_strats = _get_strats_from_topic(topics)
    monitor_id_map = defaultdict(list)
    for strat in old_strats:
      monitor_id = self._strat_monitor_id_dict[strat]
      monitor_id_map[monitor_id].append(strat)
    for monitor_id, strats in monitor_id_map.items():
      logging.info('Reuse existing monitor process for old strats: %s' % strats)
      self._launch_sub_monitor(monitor_id, self._monitor_info[monitor_id]['strats'])

  async def _check_new_topics(self):
    logging.info('Check new topics')
    query_func = functools.partial(
        _get_topics,
        self._strategy_config_filename, self._kafka_config_filename,
        topic_type=self._topic_type, business_unit=self._business_unit)
    with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor:
      all_topics = await asyncio.wait_for(
          self._ioloop.run_in_executor(executor, query_func), 10)
    new_topics = set(all_topics).difference(self._topics)
    new_strat_topics = set([
        topic for topic in new_topics if parse_kafka_topic(topic)[1].strategy_name
        not in self._strat_monitor_id_dict])
    old_strat_topics = new_topics.difference(new_strat_topics)
    self._handle_old_strat_topic(old_strat_topics)
    self._handle_new_strat_topic(new_strat_topics)
    self._topics = all_topics
    self._strats = _get_strats_from_topic(self._topics)

  def _check_running_status(self):
    if self._exit:
      return
    logging.info('Check running status')
    exit_monitor_ids = [
        monitor_id for monitor_id, info in self._monitor_info.items()
        if info['process'].poll() is not None]
    for monitor_id in exit_monitor_ids:
      strats = self._monitor_info[monitor_id]['strats']
      logging.info('Relaunch existing monitor process due to exit: %s' % strats)
      self._launch_sub_monitor(monitor_id, strats)

  def _check_exit(self):
    logging.info('Check exit')
    subprocess_exit = all([info['process'].poll() is not None for info in
                           self._monitor_info.values()])
    if subprocess_exit:
      self._exit = True
      logging.info('Exit monitor runner due to all subprocess exit.')
    if (self._exit_after_min is not None and
        get_timestamp() - self._start_ts > self._exit_after_min * 60 * 10**9):
      self._exit = True
      logging.info(
          'Exit monitor runner due to exit after %s min.' % self._exit_after_min)
    td_str = _to_td_str(get_timestamp())
    if td_str != self._td_str:
      self._exit = True
      logging.info('Exit monitor runner due to trading date change.')
    if self._exit:
      self._on_exit()

  def _on_exit(self):
    self._clean_up()
    sys.exit(0)

  def _clean_up(self):
    for monitor_id, info in self._monitor_info.items():
      process = info['process']
      strat_file = info['strat_file']
      logging.info(f'terminate monitor_id: {monitor_id}, process_id: {process.pid}')
      process.terminate()
      if os.path.exists(strat_file):
        os.remove(strat_file)

  def signal_handler(self, sig, frame):
    logging.warning('Caught signal: %s', sig)
    self._ioloop.add_callback_from_signal(self._on_exit)


def main(argv):
  init_logging()
  business_unit = FLAGS.business_unit
  if business_unit is not None:
    business_unit = [elem.strip() for elem in business_unit.split(',')]
  ioloop = IOLoop.current()
  monitor_runner = MonitorRunner(
      ioloop=ioloop,
      strategy_config_filename=FLAGS.strategy_config_filename,
      max_workers=FLAGS.max_workers,
      coin_env=FLAGS.coin_env,
      business_unit=business_unit,
      kafka_config_filename=FLAGS.kafka_config_filename,
      kafka_topic_type=FLAGS.kafka_topic_type,
      offset_hour=FLAGS.offset_hour,
      exit_after_min=FLAGS.exit_after_min,
      dump_result_after_sec=FLAGS.dump_result_after_sec,
      dump_result_frequency=FLAGS.dump_result_frequency,
      print_result=FLAGS.print_result,
      memcached_config=FLAGS.memcached_config,
      latency_db_config=FLAGS.latency_db_config,
      plot_pnl=FLAGS.plot_pnl,
      plot_order=FLAGS.plot_order,
      plot_dir=FLAGS.plot_dir,
      pnl_adj_file=FLAGS.pnl_adj_file,
      realtime_pnl_adj=FLAGS.realtime_pnl_adj,
      strat_stat_log_dir=FLAGS.strat_stat_log_dir,
      logging_level=FLAGS.logging_level)
  signal.signal(signal.SIGTERM, monitor_runner.signal_handler)
  signal.signal(signal.SIGINT, monitor_runner.signal_handler)
  try:
    ioloop.start()
  except (KeyboardInterrupt, SystemExit):
    print()
    return


if __name__ == '__main__':
  flags.DEFINE_integer('max_workers', 24, 'Max number of workers.')

  flags.DEFINE_string('kafka_config_filename',
                      '../../coin_deploy/support_monitor/config/kafka_config.json',
                      'kafka config')

  flags.DEFINE_string('kafka_topic_type', None, '')

  flags.DEFINE_string('strategy_config_filename',
                      None,
                      'strategy_config_filename')

  flags.DEFINE_string('business_unit', None, 'filter by business_unit')

  flags.DEFINE_float('offset_hour', None, 'offset_hour')

  flags.DEFINE_string('pnl_adj_file', None, 'pnl_adj_file')

  flags.DEFINE_boolean('realtime_pnl_adj', False, 'realtime pnl adj')

  flags.DEFINE_integer('dump_result_after_sec', 60, 'dump_result_after_sec')

  flags.DEFINE_integer('dump_result_frequency', 20, 'dump_result_frequency')

  flags.DEFINE_integer('exit_after_min', None, 'exit_after_min')

  flags.DEFINE_boolean('print_result', False, 'Print result')

  flags.DEFINE_string('coin_env',
                      '~/miniconda3/envs/coin2_motion_env_migration',
                      'coin_env path relative to remote_home')

  flags.DEFINE_boolean('plot_pnl', False, 'plot pnl')

  flags.DEFINE_boolean('plot_order', False, 'plot order')

  flags.DEFINE_string('memcached_config', None, 'path of database config.')

  flags.DEFINE_string('latency_db_config', None, 'path of latency db config.')

  flags.DEFINE_string('plot_dir', None, 'plot dir')

  flags.DEFINE_string('strat_stat_log_dir', None, 'strat_stat_log_dir')

  app.run(main)
